package org.donald.strategy;

import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;

public class CustomComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm<Long> {
    /**
     * @param dataSourceNames          数据源集合
     *                                 在分库时值为所有分⽚库的集合 databaseNames：⽐如表：product_order_0/product_order_1、库ds0/ds1 等
     *                                 分表时为对应分⽚库中所有分片表的集合 tablesNames
     * @param complexKeysShardingValue 分⽚属性，包括 logicTableName 为逻辑表
     * @return
     */
    @Override
    public Collection<String> doSharding(Collection<String> dataSourceNames,
                                         ComplexKeysShardingValue<Long> complexKeysShardingValue) {
        // 得到每个分⽚健对应的值
        Collection<Long> orderIdValues = this.getShardingValue(complexKeysShardingValue, "id");
        Collection<Long> userIdValues = this.getShardingValue(complexKeysShardingValue,
                "user_id");
        List<String> shardingSuffix = new ArrayList<>();
        // 对两个分⽚健取模的⽅式
        for (Long userId : userIdValues) {
            for (Long orderId : orderIdValues) {
                String suffix = userId % 2 + "_" + orderId % 2;
                for (String databaseName : dataSourceNames) {
                    if (databaseName.endsWith(suffix)) {
                        shardingSuffix.add(databaseName);
                    }
                }
            }
        }
        return shardingSuffix;
    }

    /**
     *
     * @param shardingValues 分⽚属性，包括  logicTableName 为逻辑表， columnNameAndShardingValuesMap 存储多个分⽚健 包 括key-value
     * @param key 分⽚key，id和user_id,value:分⽚value，66和99
     * @return shardingValues 集合
     */
    private Collection<Long> getShardingValue(ComplexKeysShardingValue<Long> shardingValues,
                                              final String key) {
        Collection<Long> valueSet = new ArrayList<>();
        Map<String, Collection<Long>> columnNameAndShardingValuesMap = shardingValues.getColumnNameAndShardingValuesMap();
        if (columnNameAndShardingValuesMap.containsKey(key)) {
            valueSet.addAll(columnNameAndShardingValuesMap.get(key));
        }
        return valueSet;
    }
}
